package websocket

import (
	"errors"
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"sync"
	"time"
)

type Connection struct {
	wsConn *websocket.Conn
	//读取websocket的channel
	inChan chan []byte
	//给websocket写消息的channel
	outChan   chan []byte
	closeChan chan byte
	mutex     sync.Mutex
	//closeChan 状态
	isClosed bool
}

func initConnection(wsConn *websocket.Conn) (conn *Connection) {

	conn = &Connection{
		wsConn:    wsConn,
		inChan:    make(chan []byte, 1024),
		outChan:   make(chan []byte, 1024),
		closeChan: make(chan byte, 1),
	}

	// 启动读协程
	go conn.readLoop()

	// 启动写协程
	go conn.writeLoop()

	return conn
}

// 读取websocket数据

func (c *Connection) readMessage() (data []byte, err error) {
	select {
	case data = <-c.inChan:
		return data, nil
	case <-c.closeChan:
		err = errors.New("connection is closed")
		return data, err
	}
}

// 发送消息到websocket

func (c *Connection) writeMessage(data []byte) (err error) {
	select {
	case c.outChan <- data:
		return nil
	case <-c.closeChan:
		err = errors.New("connection is closed")
		return err
	}
}

// 关闭连接

func (c *Connection) close() {

	_ = c.wsConn.Close()

	c.mutex.Lock()
	if !c.isClosed {
		close(c.closeChan)
		c.isClosed = true
	}

	c.mutex.Unlock()

}

func (c *Connection) readLoop() {

	defer c.close()

LOOP:
	for {
		_, data, err := c.wsConn.ReadMessage()
		if err != nil {
			fmt.Println("读取数据失败", err)
			break
		}
		//如果数据量过大阻塞在这里,等待inChan有空闲的位置！
		select {
		case c.inChan <- data:
		case <-c.closeChan:
			break LOOP
		}
	}
}

func (c *Connection) writeLoop() {
	defer c.close()

	var (
		data []byte
		err  error
	)

LOOP:
	for {
		select {
		case data = <-c.outChan:
		case <-c.closeChan:
			break LOOP
		}
		if err = c.wsConn.WriteMessage(websocket.TextMessage, data); err != nil {
			break
		}
	}

}

var (
	upgrade = &websocket.Upgrader{ReadBufferSize: 1024,
		WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }}
)

func ws(w http.ResponseWriter, r *http.Request) {

	// header中添加Upgrade:websocket
	wsConn, err := upgrade.Upgrade(w, r, nil)
	if err != nil {
		fmt.Println("header中添加Upgrade:websocket失败", err)
		return
	}

	conn := initConnection(wsConn)

	// 写数据
	go func(c *Connection) {
		for {
			err := c.writeMessage([]byte("heartbeat"))
			if err != nil {
				fmt.Println(err)
				return
			}
			time.Sleep(time.Second * 3)
		}
	}(conn)

	// 接受数据
	go func(c *Connection) {
		var err error
		var message []byte
		for {
			message, err = c.readMessage()
			if err != nil {
				fmt.Println(err)
				return
			}
			fmt.Println("接收数据：", string(message))
		}

	}(conn)

}
